from cs336_basics.train_bpe import train_bpe
import os
import json


def bytes2str(bytes_in: bytes):
    return ''.join([chr(b) for b in bytes_in])

# if __name__ == '__main__':
#     vocab_dict, mergers = train_bpe('data/TinyStoriesV2-GPT4-valid.txt', 10_000, ['<|endoftext|>'])
#     os.makedirs('./tmp', exist_ok=True)
#     for k, v in vocab_dict.items():
#         vocab_dict[k] = bytes2str(v)
#     with open('./tmp/vocab.json', 'w+') as f:
#         json.dump(vocab_dict, f, indent=4)
#     with open('./tmp/merges.txt', 'w+') as f:
#         for m1, m2 in mergers:
#             f.write(f"{bytes2str(m1)} {bytes2str(m2)}\n")

import torch
from thop import profile

a = torch.randn((8, 128, 800, 64))
b = torch.randn((8, 128, 800, 64))

a = torch.einsum('hbld->hbdl', a)
res = torch.einsum('hbld,hbdq->hblq', b, a)
print(res.shape)